#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
检查并修复损坏的图片文件
"""

import os
from typing import Callable, Iterable, List, Optional
from PIL import Image
from pathlib import Path


def check_and_fix_images(
    dataset_path, file_filter: Optional[Callable[[str], bool]] = None
):
    """检查并修复损坏的图片文件"""
    corrupted_files = []
    fixed_files = []

    print(f"检查数据集: {dataset_path}")

    # 遍历所有图片文件
    for root, dirs, files in os.walk(dataset_path):
        for file in files:
            if file.lower().endswith((".jpg", ".jpeg", ".png", ".bmp")):
                file_path = os.path.join(root, file)
                if file_filter and not file_filter(file_path):
                    continue
                try:
                    # 尝试打开图片
                    with Image.open(file_path) as img:
                        img.load()
                except Exception as e:
                    print(f"发现损坏的图片: {file_path} - {e}")
                    corrupted_files.append(file_path)

                    # 尝试修复
                    try:
                        # 重新打开并保存
                        with Image.open(file_path) as img:
                            # 转换为RGB模式
                            if img.mode != "RGB":
                                img = img.convert("RGB")
                            # 保存修复后的图片
                            img.save(file_path, "JPEG", quality=95)
                        print(f"已修复: {file_path}")
                        fixed_files.append(file_path)
                    except Exception as fix_e:
                        print(f"无法修复: {file_path} - {fix_e}")

    print(f"\n检查完成:")
    print(f"发现损坏文件: {len(corrupted_files)} 个")
    print(f"成功修复: {len(fixed_files)} 个")

    return corrupted_files, fixed_files


def collect_voc_datasets(
    voc_data_path: str,
    include_keywords: Optional[Iterable[str]] = None,
    exclude_keywords: Optional[Iterable[str]] = None,
) -> List[str]:
    include_keywords = [kw.lower() for kw in (include_keywords or []) if kw]
    exclude_keywords = [kw.lower() for kw in (exclude_keywords or []) if kw]

    dataset_paths: List[str] = []

    for folder in os.listdir(voc_data_path):
        folder_path = os.path.join(voc_data_path, folder)
        if not os.path.isdir(folder_path):
            continue

        folder_lower = folder.lower()
        if include_keywords and not any(kw in folder_lower for kw in include_keywords):
            continue
        if exclude_keywords and any(kw in folder_lower for kw in exclude_keywords):
            continue

        dataset_paths.append(os.path.join(folder_path, "JPEGImages"))

    return dataset_paths


if __name__ == "__main__":
    # 检查VOC数据集
    voc_datasets = []
    # voc_datasets = [
    #     "F:/GitProject/yolov7-pytorch/VOCdevkit/251003voc/bd45_voc/JPEGImages",
    #     "F:/GitProject/yolov7-pytorch/VOCdevkit/251003voc/bd46_voc/JPEGImages",
    #     "F:/GitProject/yolov7-pytorch/VOCdevkit/251003voc/ht92_voc/JPEGImages",
    #     "F:/GitProject/yolov7-pytorch/VOCdevkit/251003voc/ht93_voc/JPEGImages",
    #     "F:/GitProject/yolov7-pytorch/VOCdevkit/251003voc/ht94_voc/JPEGImages",
    #     "F:/GitProject/yolov7-pytorch/VOCdevkit/251003voc/ht95_voc/JPEGImages",
    #     "F:/GitProject/yolov7-pytorch/VOCdevkit/251003voc/ht96_voc/JPEGImages",
    #     "F:/GitProject/yolov7-pytorch/VOCdevkit/251003voc/uv68_voc/JPEGImages",
    #     "F:/GitProject/yolov7-pytorch/VOCdevkit/251003voc/uv69_voc/JPEGImages",
    #     "F:/GitProject/yolov7-pytorch/VOCdevkit/251003voc/uv71_voc/JPEGImages",
    #     "F:/GitProject/yolov7-pytorch/VOCdevkit/251003voc/uv72_voc/JPEGImages",
    #     "F:/GitProject/yolov7-pytorch/VOCdevkit/251003voc/uv75_voc/JPEGImages",
    # ]
    # 获取voc_datasets中的所有文件夹
    voc_data_path = r"F:/GitProject/yolov7-pytorch/VOCdevkit/251003voc"
    include_keywords: List[str] = []
    exclude_keywords: List[str] = []
    voc_datasets = collect_voc_datasets(
        voc_data_path,
        include_keywords=include_keywords,
        exclude_keywords=exclude_keywords,
    )

    def default_file_filter(_path: str) -> bool:
        return True

    all_corrupted = []
    all_fixed = []

    for dataset in voc_datasets:
        if os.path.exists(dataset):
            corrupted, fixed = check_and_fix_images(
                dataset, file_filter=default_file_filter
            )
            all_corrupted.extend(corrupted)
            all_fixed.extend(fixed)

    print(f"\n总体统计:")
    print(f"总损坏文件: {len(all_corrupted)} 个")
    print(f"总修复文件: {len(all_fixed)} 个")
